#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RAG Tool工具
封装RAG系统为Tool，供Agent调用
"""

import json
import logging
from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from abc import ABC, abstractmethod

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class ToolResult:
    """工具执行结果"""
    success: bool
    data: Any
    message: str
    metadata: Dict[str, Any] = None

class RAGTool(ABC):
    """RAG工具基类"""
    
    def __init__(self, name: str, description: str):
        """
        初始化RAG工具
        
        Args:
            name: 工具名称
            description: 工具描述
        """
        self.name = name
        self.description = description
    
    @abstractmethod
    def execute(self, **kwargs) -> ToolResult:
        """
        执行工具
        
        Args:
            **kwargs: 工具参数
            
        Returns:
            ToolResult: 执行结果
        """
        pass
    
    def get_schema(self) -> Dict[str, Any]:
        """
        获取工具模式
        
        Returns:
            Dict[str, Any]: 工具模式定义
        """
        return {
            "name": self.name,
            "description": self.description,
            "parameters": self._get_parameters()
        }
    
    @abstractmethod
    def _get_parameters(self) -> Dict[str, Any]:
        """
        获取参数定义
        
        Returns:
            Dict[str, Any]: 参数定义
        """
        pass

class KnowledgeBaseSearchTool(RAGTool):
    """知识库搜索工具"""
    
    def __init__(self):
        super().__init__(
            name="knowledge_base_search",
            description="搜索知识库中的实体、关系和知识块"
        )
        
        # 导入知识库搜索模块
        import sys
        import os
        sys.path.append(os.path.dirname(__file__))
        from knowledge_base_search import KnowledgeBaseSearch
        
        self.kb_search = KnowledgeBaseSearch()
    
    def execute(self, query: str, search_type: str = "all", top_k: int = 5) -> ToolResult:
        """
        执行知识库搜索
        
        Args:
            query: 搜索查询
            search_type: 搜索类型 (all, entities, relations, chunks)
            top_k: 返回结果数量
            
        Returns:
            ToolResult: 搜索结果
        """
        try:
            logger.info(f"执行知识库搜索: {query}")
            
            if search_type == "entities":
                results = self.kb_search.search_entities(query)[:top_k]
            elif search_type == "relations":
                results = self.kb_search.search_relations(query)[:top_k]
            elif search_type == "chunks":
                results = self.kb_search.search_chunks(query, top_k)
            else:  # all
                results = self.kb_search.search_all(query, top_k)
            
            # 格式化结果
            if search_type == "all":
                formatted_results = {
                    "entities": [
                        {
                            "id": result.id,
                            "title": result.title,
                            "content": result.content,
                            "score": result.score,
                            "metadata": result.metadata
                        }
                        for result in results["entities"]
                    ],
                    "relations": [
                        {
                            "id": result.id,
                            "title": result.title,
                            "content": result.content,
                            "score": result.score,
                            "metadata": result.metadata
                        }
                        for result in results["relations"]
                    ],
                    "chunks": [
                        {
                            "id": result.id,
                            "title": result.title,
                            "content": result.content,
                            "score": result.score,
                            "metadata": result.metadata
                        }
                        for result in results["chunks"]
                    ]
                }
            else:
                formatted_results = [
                    {
                        "id": result.id,
                        "title": result.title,
                        "content": result.content,
                        "score": result.score,
                        "metadata": result.metadata
                    }
                    for result in results
                ]
            
            return ToolResult(
                success=True,
                data=formatted_results,
                message=f"搜索完成，找到 {len(formatted_results) if search_type != 'all' else sum(len(v) for v in formatted_results.values())} 个结果",
                metadata={
                    "query": query,
                    "search_type": search_type,
                    "top_k": top_k
                }
            )
            
        except Exception as e:
            logger.error(f"知识库搜索失败: {e}")
            return ToolResult(
                success=False,
                data=None,
                message=f"搜索失败: {str(e)}"
            )
    
    def _get_parameters(self) -> Dict[str, Any]:
        """获取参数定义"""
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "搜索查询文本"
                },
                "search_type": {
                    "type": "string",
                    "enum": ["all", "entities", "relations", "chunks"],
                    "description": "搜索类型",
                    "default": "all"
                },
                "top_k": {
                    "type": "integer",
                    "description": "返回结果数量",
                    "default": 5,
                    "minimum": 1,
                    "maximum": 20
                }
            },
            "required": ["query"]
        }

class HybridRetrievalTool(RAGTool):
    """混合检索工具"""
    
    def __init__(self):
        super().__init__(
            name="hybrid_retrieval",
            description="执行混合检索（向量检索+关键词检索）"
        )
        
        # 导入混合检索模块
        import sys
        import os
        sys.path.append(os.path.dirname(__file__))
        from hybrid_retrieval import HybridRetrieval
        
        self.hybrid_retrieval = HybridRetrieval()
    
    def execute(self, query: str, top_k: int = 5) -> ToolResult:
        """
        执行混合检索
        
        Args:
            query: 检索查询
            top_k: 返回结果数量
            
        Returns:
            ToolResult: 检索结果
        """
        try:
            logger.info(f"执行混合检索: {query}")
            
            results = self.hybrid_retrieval.hybrid_search(query, top_k)
            
            formatted_results = [
                {
                    "id": result["id"],
                    "text": result["text"],
                    "score": result["score"],
                    "metadata": result["metadata"],
                    "retrieval_types": result["retrieval_types"],
                    "vector_score": result["vector_score"],
                    "keyword_score": result["keyword_score"]
                }
                for result in results
            ]
            
            return ToolResult(
                success=True,
                data=formatted_results,
                message=f"混合检索完成，找到 {len(formatted_results)} 个结果",
                metadata={
                    "query": query,
                    "top_k": top_k
                }
            )
            
        except Exception as e:
            logger.error(f"混合检索失败: {e}")
            return ToolResult(
                success=False,
                data=None,
                message=f"检索失败: {str(e)}"
            )
    
    def _get_parameters(self) -> Dict[str, Any]:
        """获取参数定义"""
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "检索查询文本"
                },
                "top_k": {
                    "type": "integer",
                    "description": "返回结果数量",
                    "default": 5,
                    "minimum": 1,
                    "maximum": 20
                }
            },
            "required": ["query"]
        }

class RealtimeQATool(RAGTool):
    """实时问答工具"""
    
    def __init__(self):
        super().__init__(
            name="realtime_qa",
            description="基于RAG的实时问答"
        )
        
        # 导入实时处理模块
        import sys
        import os
        sys.path.append(os.path.dirname(__file__))
        from realtime_processor import MockLLMProcessor
        
        self.qa_processor = MockLLMProcessor()
    
    def execute(self, question: str, top_k: int = 5) -> ToolResult:
        """
        执行实时问答
        
        Args:
            question: 用户问题
            top_k: 检索结果数量
            
        Returns:
            ToolResult: 问答结果
        """
        try:
            logger.info(f"执行实时问答: {question}")
            
            result = self.qa_processor.process_question(question, top_k)
            
            if result['status'] == 'success':
                return ToolResult(
                    success=True,
                    data={
                        "question": result['question'],
                        "answer": result['answer'],
                        "context_docs": result['context_docs'],
                        "retrieval_count": result['retrieval_count']
                    },
                    message="问答完成",
                    metadata={
                        "question": question,
                        "top_k": top_k,
                        "model": result.get('model', 'mock-llm')
                    }
                )
            else:
                return ToolResult(
                    success=False,
                    data=None,
                    message=result['message']
                )
            
        except Exception as e:
            logger.error(f"实时问答失败: {e}")
            return ToolResult(
                success=False,
                data=None,
                message=f"问答失败: {str(e)}"
            )
    
    def _get_parameters(self) -> Dict[str, Any]:
        """获取参数定义"""
        return {
            "type": "object",
            "properties": {
                "question": {
                    "type": "string",
                    "description": "用户问题"
                },
                "top_k": {
                    "type": "integer",
                    "description": "检索结果数量",
                    "default": 5,
                    "minimum": 1,
                    "maximum": 20
                }
            },
            "required": ["question"]
        }

class KnowledgeGraphTool(RAGTool):
    """知识图谱工具"""
    
    def __init__(self):
        super().__init__(
            name="knowledge_graph",
            description="构建和查询知识图谱"
        )
        
        # 导入知识库搜索模块
        import sys
        import os
        sys.path.append(os.path.dirname(__file__))
        from knowledge_base_search import KnowledgeBaseSearch
        
        self.kb_search = KnowledgeBaseSearch()
    
    def execute(self, entity_name: str, depth: int = 2) -> ToolResult:
        """
        构建知识图谱
        
        Args:
            entity_name: 中心实体名称
            depth: 图谱深度
            
        Returns:
            ToolResult: 知识图谱结果
        """
        try:
            logger.info(f"构建知识图谱: {entity_name}, 深度: {depth}")
            
            graph = self.kb_search.get_knowledge_graph(entity_name, depth)
            
            return ToolResult(
                success=True,
                data=graph,
                message=f"知识图谱构建完成，节点数: {len(graph['nodes'])}, 边数: {len(graph['edges'])}",
                metadata={
                    "center_entity": entity_name,
                    "depth": depth
                }
            )
            
        except Exception as e:
            logger.error(f"构建知识图谱失败: {e}")
            return ToolResult(
                success=False,
                data=None,
                message=f"构建失败: {str(e)}"
            )
    
    def _get_parameters(self) -> Dict[str, Any]:
        """获取参数定义"""
        return {
            "type": "object",
            "properties": {
                "entity_name": {
                    "type": "string",
                    "description": "中心实体名称"
                },
                "depth": {
                    "type": "integer",
                    "description": "图谱深度",
                    "default": 2,
                    "minimum": 1,
                    "maximum": 5
                }
            },
            "required": ["entity_name"]
        }

class RAGToolManager:
    """RAG工具管理器"""
    
    def __init__(self):
        """初始化工具管理器"""
        self.tools = {}
        self._register_tools()
    
    def _register_tools(self):
        """注册所有RAG工具"""
        self.tools = {
            "knowledge_base_search": KnowledgeBaseSearchTool(),
            "hybrid_retrieval": HybridRetrievalTool(),
            "realtime_qa": RealtimeQATool(),
            "knowledge_graph": KnowledgeGraphTool()
        }
        
        logger.info(f"已注册 {len(self.tools)} 个RAG工具")
    
    def get_tool(self, tool_name: str) -> Optional[RAGTool]:
        """
        获取工具
        
        Args:
            tool_name: 工具名称
            
        Returns:
            Optional[RAGTool]: 工具实例
        """
        return self.tools.get(tool_name)
    
    def list_tools(self) -> List[Dict[str, Any]]:
        """
        列出所有工具
        
        Returns:
            List[Dict[str, Any]]: 工具列表
        """
        return [tool.get_schema() for tool in self.tools.values()]
    
    def execute_tool(self, tool_name: str, **kwargs) -> ToolResult:
        """
        执行工具
        
        Args:
            tool_name: 工具名称
            **kwargs: 工具参数
            
        Returns:
            ToolResult: 执行结果
        """
        tool = self.get_tool(tool_name)
        if not tool:
            return ToolResult(
                success=False,
                data=None,
                message=f"工具不存在: {tool_name}"
            )
        
        return tool.execute(**kwargs)
    
    def get_tools_schema(self) -> Dict[str, Any]:
        """
        获取所有工具的Schema
        
        Returns:
            Dict[str, Any]: 工具Schema
        """
        return {
            "tools": self.list_tools(),
            "total_tools": len(self.tools)
        }


def main():
    """测试RAG工具"""
    print("测试RAG工具功能")
    print("=" * 50)
    
    # 创建工具管理器
    tool_manager = RAGToolManager()
    
    # 列出所有工具
    print("可用工具:")
    for tool in tool_manager.list_tools():
        print(f"  - {tool['name']}: {tool['description']}")
    
    # 测试知识库搜索工具
    print("\n测试知识库搜索工具:")
    result = tool_manager.execute_tool("knowledge_base_search", query="人与自然", top_k=3)
    if result.success:
        print(f"搜索成功: {result.message}")
        print(f"结果数量: {len(result.data['entities']) + len(result.data['chunks'])}")
    else:
        print(f"搜索失败: {result.message}")
    
    # 测试混合检索工具
    print("\n测试混合检索工具:")
    result = tool_manager.execute_tool("hybrid_retrieval", query="生态危机", top_k=3)
    if result.success:
        print(f"检索成功: {result.message}")
        print(f"结果数量: {len(result.data)}")
    else:
        print(f"检索失败: {result.message}")
    
    # 测试实时问答工具
    print("\n测试实时问答工具:")
    result = tool_manager.execute_tool("realtime_qa", question="什么是生态危机？", top_k=3)
    if result.success:
        print(f"问答成功: {result.message}")
        print(f"回答: {result.data['answer'][:100]}...")
    else:
        print(f"问答失败: {result.message}")
    
    # 测试知识图谱工具
    print("\n测试知识图谱工具:")
    result = tool_manager.execute_tool("knowledge_graph", entity_name="人与自然", depth=2)
    if result.success:
        print(f"图谱构建成功: {result.message}")
        print(f"节点数: {len(result.data['nodes'])}")
        print(f"边数: {len(result.data['edges'])}")
    else:
        print(f"图谱构建失败: {result.message}")


if __name__ == "__main__":
    main()


